addpath ../
addpath ../export_fig/


% ****************************************************************
% test imputation methods
% 
ntis = 53;
ngene = 100;
nind = 450;
nfactor = 15;
factor_per_gene = 3;
pmiss = .8;

U_true = randn(ntis,nfactor,'single');
Y_full = zeros(ntis,nind,ngene,'single');
Y_train = zeros(ntis,nind,ngene,'single');

for g = 1:ngene,

    V = randn(factor_per_gene,nind,'single');
    yy = zscore(U_true(:,randsample(nfactor,factor_per_gene))*V,[],2);
    Y_full(:,:,g) = yy;

    yy(rand(ntis,nind,'single') <= pmiss) = nan;
    Y_train(:,:,g) = yy;
    clear yy;
end

% ****************************************************************
J = 30;
burnin = 10;
inittime = 200;
nepoch = 500;
batchsize = 50;

outdir = 'sim04_result';
! rm -rf sim04_result
run_modQTL_imputation_gibbs(Y_train,J,burnin,inittime,nepoch,batchsize,outdir);
% run_modQTL_imputation_memo(Y_train,J,burnin,nepoch,batchsize,outdir);


% ****************************************************************
% 1. measure correlation between true vs imputed
load([outdir,'/parameters.mat']);

files = dir([outdir,'/V_minibatch*.mat']);
yytrue = [];
yyhat = [];

for ff = 1:numel(files),
    load([outdir,'/',files(ff).name]);
    
    Y = Y_full(:,:,genes_minibatch);
    missing = isnan(Y_train(:,:,genes_minibatch));

    yytrue = [yytrue; Y(missing)];

    YY = zeros(ntis,nind,numel(genes_minibatch));

    for g = 1:numel(genes_minibatch),
        YY(:,:,g) = U * V_minibatch(:,:,g);
    end

    yyhat = [yyhat; YY(missing)];

    clear Y *_minibatch;
end

[r,p] = corr(yytrue,yyhat);

h1 = figure(1);
set(h1,'position',[.5 .5 400 300]);
idx = randsample(numel(yytrue),10000);
hold on;
dscatter( yytrue(idx), yyhat(idx) );
plot( [-3,3], r*[-3,3], 'r--', 'linewidth', 2 );
plot( [-3,3], [-3,3], 'k--', 'linewidth', 2 );
hold off;
xlabel('true'); ylabel('imputed');
title(sprintf('r = %.2e p = %.2e',r,p));
axis('tight');

export_fig('sim04_fig01','-pdf','-a1','-m2','-transparent');
